package io.debezium.connector.postgresql;

import com.huawei.gaussdb.jdbc.core.BaseConnection;
import com.huawei.gaussdb.jdbc.core.TypeInfo;
import io.debezium.DebeziumException;
import io.debezium.annotation.Immutable;
import io.debezium.connector.postgresql.connection.GaussDBConnection;
import io.debezium.util.Collect;
import org.apache.kafka.connect.errors.ConnectException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

public class TypeRegistry {
    private static final Logger LOGGER = LoggerFactory.getLogger(TypeRegistry.class);
    public static final String TYPE_NAME_GEOGRAPHY = "geography";
    public static final String TYPE_NAME_GEOMETRY = "geometry";
    public static final String TYPE_NAME_CITEXT = "citext";
    public static final String TYPE_NAME_HSTORE = "hstore";
    public static final String TYPE_NAME_LTREE = "ltree";
    public static final String TYPE_NAME_HSTORE_ARRAY = "_hstore";
    public static final String TYPE_NAME_GEOGRAPHY_ARRAY = "_geography";
    public static final String TYPE_NAME_GEOMETRY_ARRAY = "_geometry";
    public static final String TYPE_NAME_CITEXT_ARRAY = "_citext";
    public static final String TYPE_NAME_LTREE_ARRAY = "_ltree";
    public static final int NO_TYPE_MODIFIER = -1;
    public static final int UNKNOWN_LENGTH = -1;
    public static final int DOMAIN_TYPE = 2001;
    private static final String CATEGORY_ARRAY = "A";
    private static final String CATEGORY_ENUM = "E";
    private static final String SQL_ENUM_VALUES =
            "SELECT t.enumtypid as id, array_agg(t.enumlabel) as values FROM pg_catalog.pg_enum t GROUP BY id";
    private static final String SQL_TYPES =
            "SELECT t.oid AS oid, t.typname AS name, t.typelem AS element, t.typbasetype AS parentoid, t.typtypmod as modifiers, t.typcategory as category, e.values as enum_values FROM pg_catalog.pg_type t JOIN pg_catalog.pg_namespace n ON (t.typnamespace = n.oid) LEFT JOIN (SELECT t.enumtypid as id, array_agg(t.enumlabel) as values FROM pg_catalog.pg_enum t GROUP BY id) e ON (t.oid = e.id) WHERE n.nspname != 'pg_toast'";
    private static final String SQL_NAME_LOOKUP =
            "SELECT t.oid AS oid, t.typname AS name, t.typelem AS element, t.typbasetype AS parentoid, t.typtypmod as modifiers, t.typcategory as category, e.values as enum_values FROM pg_catalog.pg_type t JOIN pg_catalog.pg_namespace n ON (t.typnamespace = n.oid) LEFT JOIN (SELECT t.enumtypid as id, array_agg(t.enumlabel) as values FROM pg_catalog.pg_enum t GROUP BY id) e ON (t.oid = e.id) WHERE n.nspname != 'pg_toast' AND t.typname = ?";
    private static final String SQL_OID_LOOKUP =
            "SELECT t.oid AS oid, t.typname AS name, t.typelem AS element, t.typbasetype AS parentoid, t.typtypmod as modifiers, t.typcategory as category, e.values as enum_values FROM pg_catalog.pg_type t JOIN pg_catalog.pg_namespace n ON (t.typnamespace = n.oid) LEFT JOIN (SELECT t.enumtypid as id, array_agg(t.enumlabel) as values FROM pg_catalog.pg_enum t GROUP BY id) e ON (t.oid = e.id) WHERE n.nspname != 'pg_toast' AND t.oid = ?";
    private static final Map<String, String> LONG_TYPE_NAMES =
            Collections.unmodifiableMap(getLongTypeNames());
    private final Map<String, PostgresType> nameToType = new HashMap();
    private final Map<Integer, PostgresType> oidToType = new HashMap();
    private final GaussDBConnection connection;
    private final SqlTypeMapper sqlTypeMapper;
    private int geometryOid = Integer.MIN_VALUE;
    private int geographyOid = Integer.MIN_VALUE;
    private int citextOid = Integer.MIN_VALUE;
    private int hstoreOid = Integer.MIN_VALUE;
    private int ltreeOid = Integer.MIN_VALUE;
    private int hstoreArrayOid = Integer.MIN_VALUE;
    private int geometryArrayOid = Integer.MIN_VALUE;
    private int geographyArrayOid = Integer.MIN_VALUE;
    private int citextArrayOid = Integer.MIN_VALUE;
    private int ltreeArrayOid = Integer.MIN_VALUE;

    private static Map<String, String> getLongTypeNames() {
        Map<String, String> longTypeNames = new HashMap();
        longTypeNames.put("bigint", "int8");
        longTypeNames.put("bit varying", "varbit");
        longTypeNames.put("boolean", "bool");
        longTypeNames.put("character", "bpchar");
        longTypeNames.put("character varying", "varchar");
        longTypeNames.put("double precision", "float8");
        longTypeNames.put("integer", "int4");
        longTypeNames.put("real", "float4");
        longTypeNames.put("smallint", "int2");
        longTypeNames.put("timestamp without time zone", "timestamp");
        longTypeNames.put("timestamp with time zone", "timestamptz");
        longTypeNames.put("time without time zone", "time");
        longTypeNames.put("time with time zone", "timetz");
        return longTypeNames;
    }

    public TypeRegistry(GaussDBConnection connection) {
        try {
            this.connection = connection;
            this.sqlTypeMapper = new SqlTypeMapper(this.connection);
            this.prime();
        } catch (SQLException var3) {
            SQLException e = var3;
            throw new DebeziumException("Couldn't initialize type registry", e);
        }
    }

    private void addType(PostgresType type) {
        this.oidToType.put(type.getOid(), type);
        this.nameToType.put(type.getName(), type);
        if ("geometry".equals(type.getName())) {
            this.geometryOid = type.getOid();
        } else if ("geography".equals(type.getName())) {
            this.geographyOid = type.getOid();
        } else if ("citext".equals(type.getName())) {
            this.citextOid = type.getOid();
        } else if ("hstore".equals(type.getName())) {
            this.hstoreOid = type.getOid();
        } else if ("ltree".equals(type.getName())) {
            this.ltreeOid = type.getOid();
        } else if ("_hstore".equals(type.getName())) {
            this.hstoreArrayOid = type.getOid();
        } else if ("_geometry".equals(type.getName())) {
            this.geometryArrayOid = type.getOid();
        } else if ("_geography".equals(type.getName())) {
            this.geographyArrayOid = type.getOid();
        } else if ("_citext".equals(type.getName())) {
            this.citextArrayOid = type.getOid();
        } else if ("_ltree".equals(type.getName())) {
            this.ltreeArrayOid = type.getOid();
        }
    }

    public PostgresType get(int oid) {
        PostgresType r = (PostgresType) this.oidToType.get(oid);
        if (r == null) {
            r = this.resolveUnknownType(oid);
            if (r == null) {
                LOGGER.warn("Unknown OID {} requested", oid);
                r = PostgresType.UNKNOWN;
            }
        }

        return r;
    }

    public PostgresType get(String name) {
        switch (name) {
            case "serial":
                name = "int4";
                break;
            case "smallserial":
                name = "int2";
                break;
            case "bigserial":
                name = "int8";
        }

        String[] parts = name.split("\\.");
        if (parts.length > 1) {
            name = parts[1];
        }

        if (name.charAt(0) == '"') {
            name = name.substring(1, name.length() - 1);
        }

        PostgresType r = (PostgresType) this.nameToType.get(name);
        if (r == null) {
            r = this.resolveUnknownType(name);
            if (r == null) {
                LOGGER.warn("Unknown type named {} requested", name);
                r = PostgresType.UNKNOWN;
            }
        }

        return r;
    }

    public int geometryOid() {
        return this.geometryOid;
    }

    public int geographyOid() {
        return this.geographyOid;
    }

    public int citextOid() {
        return this.citextOid;
    }

    public int hstoreOid() {
        return this.hstoreOid;
    }

    public int ltreeOid() {
        return this.ltreeOid;
    }

    public int hstoreArrayOid() {
        return this.hstoreArrayOid;
    }

    public int geometryArrayOid() {
        return this.geometryArrayOid;
    }

    public int geographyArrayOid() {
        return this.geographyArrayOid;
    }

    public int citextArrayOid() {
        return this.citextArrayOid;
    }

    public int ltreeArrayOid() {
        return this.ltreeArrayOid;
    }

    public static String normalizeTypeName(String typeName) {
        return (String) LONG_TYPE_NAMES.getOrDefault(typeName, typeName);
    }

    private void prime() throws SQLException {
        Statement statement = this.connection.connection().createStatement();

        try {
            ResultSet rs =
                    statement.executeQuery(
                            "SELECT t.oid AS oid, t.typname AS name, t.typelem AS element, t.typbasetype AS parentoid, t.typtypmod as modifiers, t.typcategory as category, e.values as enum_values FROM pg_catalog.pg_type t JOIN pg_catalog.pg_namespace n ON (t.typnamespace = n.oid) LEFT JOIN (SELECT t.enumtypid as id, array_agg(t.enumlabel) as values FROM pg_catalog.pg_enum t GROUP BY id) e ON (t.oid = e.id) WHERE n.nspname != 'pg_toast'");

            try {
                List<PostgresType.Builder> delayResolvedBuilders = new ArrayList();

                while (rs.next()) {
                    PostgresType.Builder builder = this.createTypeBuilderFromResultSet(rs);
                    if (!builder.hasParentType()) {
                        this.addType(builder.build());
                    } else {
                        delayResolvedBuilders.add(builder);
                    }
                }

                Iterator var10 = delayResolvedBuilders.iterator();

                while (var10.hasNext()) {
                    PostgresType.Builder builder = (PostgresType.Builder) var10.next();
                    this.addType(builder.build());
                }
            } catch (Throwable var8) {
                if (rs != null) {
                    try {
                        rs.close();
                    } catch (Throwable var7) {
                        var8.addSuppressed(var7);
                    }
                }

                throw var8;
            }

            if (rs != null) {
                rs.close();
            }
        } catch (Throwable var9) {
            if (statement != null) {
                try {
                    statement.close();
                } catch (Throwable var6) {
                    var9.addSuppressed(var6);
                }
            }

            throw var9;
        }

        if (statement != null) {
            statement.close();
        }
    }

    private PostgresType.Builder createTypeBuilderFromResultSet(ResultSet rs) throws SQLException {
        int oid = (int) rs.getLong("oid");
        int parentTypeOid = (int) rs.getLong("parentoid");
        int modifiers = (int) rs.getLong("modifiers");
        String typeName = rs.getString("name");
        String category = rs.getString("category");
        PostgresType.Builder builder =
                new PostgresType.Builder(
                        this,
                        typeName,
                        oid,
                        this.sqlTypeMapper.getSqlType(typeName),
                        modifiers,
                        getTypeInfo(this.connection));
        if ("E".equals(category)) {
            String[] enumValues = (String[]) rs.getArray("enum_values").getArray();
            builder = builder.enumValues(Arrays.asList(enumValues));
        } else if ("A".equals(category)) {
            builder = builder.elementType((int) rs.getLong("element"));
        }

        return builder.parentType(parentTypeOid);
    }

    private PostgresType resolveUnknownType(String name) {
        try {
            LOGGER.trace("Type '{}' not cached, attempting to lookup from database.", name);
            PreparedStatement statement =
                    this.connection
                            .connection()
                            .prepareStatement(
                                    "SELECT t.oid AS oid, t.typname AS name, t.typelem AS element, t.typbasetype AS parentoid, t.typtypmod as modifiers, t.typcategory as category, e.values as enum_values FROM pg_catalog.pg_type t JOIN pg_catalog.pg_namespace n ON (t.typnamespace = n.oid) LEFT JOIN (SELECT t.enumtypid as id, array_agg(t.enumlabel) as values FROM pg_catalog.pg_enum t GROUP BY id) e ON (t.oid = e.id) WHERE n.nspname != 'pg_toast' AND t.typname = ?");

            PostgresType var3;
            try {
                statement.setString(1, name);
                var3 = this.loadType(statement);
            } catch (Throwable var6) {
                if (statement != null) {
                    try {
                        statement.close();
                    } catch (Throwable var5) {
                        var6.addSuppressed(var5);
                    }
                }

                throw var6;
            }

            if (statement != null) {
                statement.close();
            }

            return var3;
        } catch (SQLException var7) {
            SQLException e = var7;
            throw new ConnectException(
                    "Database connection failed during resolving unknown type", e);
        }
    }

    private PostgresType resolveUnknownType(int lookupOid) {
        try {
            LOGGER.trace(
                    "Type OID '{}' not cached, attempting to lookup from database.", lookupOid);
            PreparedStatement statement =
                    this.connection
                            .connection()
                            .prepareStatement(
                                    "SELECT t.oid AS oid, t.typname AS name, t.typelem AS element, t.typbasetype AS parentoid, t.typtypmod as modifiers, t.typcategory as category, e.values as enum_values FROM pg_catalog.pg_type t JOIN pg_catalog.pg_namespace n ON (t.typnamespace = n.oid) LEFT JOIN (SELECT t.enumtypid as id, array_agg(t.enumlabel) as values FROM pg_catalog.pg_enum t GROUP BY id) e ON (t.oid = e.id) WHERE n.nspname != 'pg_toast' AND t.oid = ?");

            PostgresType var3;
            try {
                statement.setInt(1, lookupOid);
                var3 = this.loadType(statement);
            } catch (Throwable var6) {
                if (statement != null) {
                    try {
                        statement.close();
                    } catch (Throwable var5) {
                        var6.addSuppressed(var5);
                    }
                }

                throw var6;
            }

            if (statement != null) {
                statement.close();
            }

            return var3;
        } catch (SQLException var7) {
            SQLException e = var7;
            throw new ConnectException(
                    "Database connection failed during resolving unknown type", e);
        }
    }

    private PostgresType loadType(PreparedStatement statement) throws SQLException {
        ResultSet rs = statement.executeQuery();

        PostgresType var4;
        label40:
        {
            try {
                if (rs.next()) {
                    PostgresType result = this.createTypeBuilderFromResultSet(rs).build();
                    this.addType(result);
                    var4 = result;
                    break label40;
                }
            } catch (Throwable var6) {
                if (rs != null) {
                    try {
                        rs.close();
                    } catch (Throwable var5) {
                        var6.addSuppressed(var5);
                    }
                }

                throw var6;
            }

            if (rs != null) {
                rs.close();
            }

            return null;
        }

        if (rs != null) {
            rs.close();
        }

        return var4;
    }

    private static TypeInfo getTypeInfo(GaussDBConnection connection) throws SQLException {
        return ((BaseConnection) connection.connection()).getTypeInfo();
    }

    private static class SqlTypeMapper {
        private static final String SQL_TYPE_DETAILS =
                "SELECT DISTINCT ON (typname) typname, typinput='array_in'::regproc, typtype, sp.r, pg_type.oid   FROM pg_catalog.pg_type   LEFT   JOIN (select ns.oid as nspoid, ns.nspname, r.r           from pg_namespace as ns           join ( select s.r, (current_schemas(false))[s.r] as nspname                    from generate_series(1, array_upper(current_schemas(false), 1)) as s(r) ) as r          using ( nspname )        ) as sp     ON sp.nspoid = typnamespace  ORDER BY typname, sp.r, pg_type.oid;";
        private final GaussDBConnection connection;
        @Immutable private final Set<String> preloadedSqlTypes;
        @Immutable private final Map<String, Integer> sqlTypesByPgTypeNames;

        private SqlTypeMapper(GaussDBConnection connection) throws SQLException {
            this.connection = connection;
            this.preloadedSqlTypes =
                    Collect.unmodifiableSet(
                            TypeRegistry.getTypeInfo(connection).getPGTypeNamesWithSQLTypes());
            this.sqlTypesByPgTypeNames = Collections.unmodifiableMap(getSqlTypes(connection));
        }

        public int getSqlType(String typeName) throws SQLException {
            boolean isCoreType = this.preloadedSqlTypes.contains(typeName);
            if (isCoreType) {
                return TypeRegistry.getTypeInfo(this.connection).getSQLType(typeName);
            } else if (typeName.endsWith("[]")) {
                return 2003;
            } else {
                try {
                    Integer pgType = (Integer) this.sqlTypesByPgTypeNames.get(typeName);
                    if (pgType != null) {
                        return pgType;
                    } else {
                        TypeRegistry.LOGGER.info(
                                "Failed to obtain SQL type information for type {} via custom statement, falling back to TypeInfo#getSQLType()",
                                typeName);
                        return TypeRegistry.getTypeInfo(this.connection).getSQLType(typeName);
                    }
                } catch (Exception var4) {
                    Exception e = var4;
                    TypeRegistry.LOGGER.warn(
                            "Failed to obtain SQL type information for type {} via custom statement, falling back to TypeInfo#getSQLType()",
                            typeName,
                            e);
                    return TypeRegistry.getTypeInfo(this.connection).getSQLType(typeName);
                }
            }
        }

        private static Map<String, Integer> getSqlTypes(GaussDBConnection connection)
                throws SQLException {
            Map<String, Integer> sqlTypesByPgTypeNames = new HashMap();
            Statement statement = connection.connection().createStatement();

            try {
                ResultSet rs =
                        statement.executeQuery(
                                "SELECT DISTINCT ON (typname) typname, typinput='array_in'::regproc, typtype, sp.r, pg_type.oid   FROM pg_catalog.pg_type   LEFT   JOIN (select ns.oid as nspoid, ns.nspname, r.r           from pg_namespace as ns           join ( select s.r, (current_schemas(false))[s.r] as nspname                    from generate_series(1, array_upper(current_schemas(false), 1)) as s(r) ) as r          using ( nspname )        ) as sp     ON sp.nspoid = typnamespace  ORDER BY typname, sp.r, pg_type.oid;");

                short type;
                try {
                    for (;
                            rs.next();
                            sqlTypesByPgTypeNames.put(rs.getString(1), Integer.valueOf(type))) {
                        boolean isArray = rs.getBoolean(2);
                        String typtype = rs.getString(3);
                        if (isArray) {
                            type = 2003;
                        } else if ("c".equals(typtype)) {
                            type = 2002;
                        } else if ("d".equals(typtype)) {
                            type = 2001;
                        } else if ("e".equals(typtype)) {
                            type = 12;
                        } else {
                            type = 1111;
                        }
                    }
                } catch (Throwable var9) {
                    if (rs != null) {
                        try {
                            rs.close();
                        } catch (Throwable var8) {
                            var9.addSuppressed(var8);
                        }
                    }

                    throw var9;
                }

                if (rs != null) {
                    rs.close();
                }
            } catch (Throwable var10) {
                if (statement != null) {
                    try {
                        statement.close();
                    } catch (Throwable var7) {
                        var10.addSuppressed(var7);
                    }
                }

                throw var10;
            }

            if (statement != null) {
                statement.close();
            }

            return sqlTypesByPgTypeNames;
        }
    }
}
